import torch
import torch.nn as nn
from transformers import BertModel, BertConfig
import torch.nn.functional as F
import os

class BertTextCNN(nn.Module):
    def __init__(self, bert_model_path, num_labels=14, filter_sizes=(2, 3, 4), num_filters=50):
        super(BertTextCNN, self).__init__()
        
        # 检查传入的路径是否为目录，如果是目录则使用本地模型
        if os.path.isdir(bert_model_path):
            # 从本地加载BERT模型
            config = BertConfig.from_json_file(os.path.join(bert_model_path, 'config.json'))
            self.bert = BertModel(config)
            # 加载预训练权重
            state_dict = torch.load(os.path.join(bert_model_path, 'pytorch_model.bin'), map_location='cpu')
            self.bert.load_state_dict(state_dict)
        else:
            # 从Hugging Face下载BERT模型
            self.bert = BertModel.from_pretrained(bert_model_path)
        
        # 冻结BERT部分参数以减少计算量
        for param in self.bert.embeddings.parameters():
            param.requires_grad = False
        
        for i in range(8):  # 冻结前8层，只微调后面的层
            for param in self.bert.encoder.layer[i].parameters():
                param.requires_grad = False

        self.convs = nn.ModuleList([nn.Conv2d(1, num_filters, (k, 256)) for k in filter_sizes])
        self.dropout = nn.Dropout(0.1)
        self.fc = nn.Linear(num_filters * len(filter_sizes), num_labels)

    def forward(self, input_ids):
        bert_output = self.bert(input_ids=input_ids)[0]  # (batch_size, seq_len, embed_dim)
        bert_output = bert_output.unsqueeze(1)  # (batch_size, 1, seq_len, embed_dim)

        conv_outputs = []
        for i, conv in enumerate(self.convs):
            conv_out = conv(bert_output).squeeze(3)  # (batch_size, num_filters, seq_len - filter_size + 1)
            pooled = F.max_pool1d(conv_out, conv_out.size(2)).squeeze(2)  # (batch_size, num_filters)
            conv_outputs.append(pooled)

        concat_output = torch.cat(conv_outputs, 1)  # (batch_size, num_filters * len(filter_sizes))
        dropout_output = self.dropout(concat_output)
        logits = self.fc(dropout_output)
        return logits
